package com.jojo.rateLimiter;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * 令牌桶算法
 *
 * @author hpl
 * @since 2024/12/16
 */
public class TokenBucket {
    // 令牌桶的最大容量
    private final int capacity;
    // 令牌生成速率（每秒生成的令牌数）
    private final int refillRate;
    // 当前桶中的令牌数量
    private AtomicInteger currentTokens = new AtomicInteger(0);
    // 上次令牌生成时间
    private long lastRefillTime;

    public TokenBucket(int capacity, int refillRate) {
        this.capacity = capacity;
        this.refillRate = refillRate;
        this.lastRefillTime = System.currentTimeMillis();
    }

    /**
     * 尝试获取一个令牌
     *
     * @return 如果成功获取到令牌，返回 true；否则返回 false
     */
    public synchronized boolean tryConsume() {
        // 进行令牌补充
        refill();

        // 如果还有令牌，则消耗一个令牌
        if (currentTokens.get() > 0) {
            currentTokens.decrementAndGet();
            return true;
        } else {
            // 如果没有令牌，请求被拒绝
            return false;
        }
    }

    /**
     * 根据时间间隔生成令牌，补充到桶中
     */
    private void refill() {
        long now = System.currentTimeMillis();
        long elapsedTime = now - lastRefillTime;

        // 计算需要补充的令牌数量
        int tokensToAdd = (int) (elapsedTime / 1000.0 * refillRate);
        if (tokensToAdd > 0) {
            // 将补充后的令牌数量更新到桶中
            currentTokens.updateAndGet(tokens -> Math.min(capacity, tokens + tokensToAdd));
            // 更新最后补充时间
            lastRefillTime = now;
        }
    }

    public static void main(String[] args) throws InterruptedException {
        // 创建一个令牌桶，最大容量为 10，每秒生成 5 个令牌
        TokenBucket bucket = new TokenBucket(10, 5);

        // 模拟请求
        for (int i = 0; i < 20; i++) {
            if (bucket.tryConsume()) {
                System.out.println("第 " + (i + 1) + " 个请求成功处理");
            } else {
                System.out.println("第 " + (i + 1) + " 个请求被拒绝");
            }
            Thread.sleep(300); // 每隔 300 毫秒发起一个请求
        }
    }
}
